from torch.utils.data import DataLoader
from utils.dataset import DigitDataset
from models.crnn import CRNN
import torch
from torch import nn
import torch.optim as optim


def train():
    dataset = DigitDataset("train_data/")
    loader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    model = CRNN(32, 1, 11, 256)  # 10数字+小数点=11类
    criterion = nn.CTCLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    for epoch in range(10):
        for inputs, targets in loader:
            seq_len = torch.IntTensor([inputs.size(0)]*inputs.size(1))
            preds = model(inputs)
            
            loss = criterion(preds.log_softmax(2), 
                           targets, 
                           seq_len, 
                           torch.IntTensor([len(t) for t in targets]))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch} Loss: {loss.item()}")

    torch.save(model.state_dict(), "crnn_model.pth")
